DLO-JZ Fully Sharded Data parallelism - Jour 4¶

Utilisation de la FSDP sur un modèle de langue Llama 3.2 3B.

Monstertruck

Objet du notebook¶

Le but de ce notebook est d'optimiser un code d'apprentissage d'un modèle Llama 3.2 sur un dataset de roleplay Imagenet :

  • Passage de DDP à FSDP
  • FSDP optimisée grâce au wrap récursif
  • Bonus : Application de la compilation PyTorch par dessus la FSDP

Les cellules dans ce notebook ne sont pas prévues pour être modifiées, sauf rares exceptions indiquées dans les commentaires. Les TP se feront en modifiant le code fsdp.py.

Les directives de modification seront marquées par l'étiquette TODO : dans le notebook suivant.

Les solutions sont présentes dans le répertoire solutions/.

Notebook rédigé par l'équipe assistance IA de l'IDRIS, octobre 2024


Environnement de calcul¶

Les fonctions python de gestion de queue Slurm dévelopées par l'IDRIS et les fonctions dédiées à la formation DLO-JZ sont à importer.

Le module d'environnement pour les jobs et la taille des images sont fixés pour ce notebook.

TODO : choisir un pseudonyme pour vous différencier dans la queue Slurm et dans les outils collaboratifs pendant la formation.

In [1]:
from idr_pytools import display_slurm_queue, gpu_jobs_submitter, search_log
from dlojz_tools import controle_technique, compare, GPU_underthehood, plot_accuracy, lrfind_plot, pipe_memory, turbo_profiler, comm_profiler
MODULE = 'pytorch-gpu/py3/2.4.0'
account = 'for@a100'
name = 'NATHOUNET'

Gestion de la queue Slurm¶

Cette partie permet d'afficher et de gérer la queue Slurm.

Pour afficher toute la queue utilisateur :

In [2]:
#display_slurm_queue(name)

Remarque: Cette fonction utilisée plusieurs fois dans ce notebook permet d'afficher la queue de manière dynamique, rafraichie toutes les 5 secondes. Cependant elle ne s'arrête que lorsque la queue est vide. Si vous désirez reprendre la main sur le notebook, il vous suffira d'arrêter manuellement la cellule avec le bouton stop. Cela a bien sûr aucun impact sur le scheduler Slurm. Les jobs ne seront pas arrêtés.

Si vous voulez arrêter des jobs dans la queue :

  • Annuler tous vos jobs dans la queue (décommenter la ligne suivante)
  • Annuler un job dans votre queue (décommenter la ligne suivante et ajouter le numéro du job à la fin de la ligne)
In [4]:
#!scancel -u $USER

Différence de scripts ¶

Pour le debug ou pour comparer son code avec les solutions mises à disposition, la fonction suivante permet d'afficher une page html contenant un différentiel de fichiers texte.

In [32]:
s1 = "./solutions/fsdp_2.py"
s2 = "./solutions/fsdp_3.py"
compare(s1, s2)

Voir le résultat du différentiel de fichiers sur la page suivante (attention au spoil !) :

compare.html


Première exécution en DDP¶

In [5]:
!cp solutions/fsdp_0.py fsdp.py


Prenez connaissance du script fsdp.py. C'est un fine-tuning d'un modèle de langue Llama 3.2 à 3 milliards de paramètres sur un dataset de Roleplay récupéré sur HuggingFace.

La structure entre ce script et celui lié à la computer vision des autres jours est très similaire.

Note : on utilise les mêmes visualisations que dans les autres TPs donc les schémas peuvent parler d'images, mais dans notre contexte de NLP, la dimension du batch ne fait pas référence à des images mais plutôt aux nombres de séquences.

In [6]:
n_gpu = 4
command = f'fsdp.py --batch-size 4 --num-workers 2 --seq-len 512 --test --nccl-profile'
jobid = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
                   account=account, time_max='00:10:00', qos="qos_gpu_a100-dev")
print(f'jobid = {jobid}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task
Submitted batch job 334759
jobid = ['334759']
In [7]:
display_slurm_queue(name)
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
            334759    gpu_p5 NATHOUNE  cfor032  R       1:32      1 jean-zay-iam31

 Done!
In [8]:
#jobid = ['334031']
In [9]:
controle_technique(jobid)
Train throughput: 28.75 images/second
GPU throughput: 28.80 images/second
epoch time: 200.37 seconds
-----------
training step time average (fwd/bkwd on GPU): 0.555491 sec (38.9%/57.5%) +/- 0.053871
loading step time average (IO + CPU to GPU transfer): 0.001095 sec +/- 0.000453

Click here to display the log file

In [10]:
turbo_profiler(jobid)
>>> Turbo Profiler >>> Training complete in 29.487146 s
No description has been provided for this image
In [11]:
comm_profiler(jobid, n_display=100)
No description has been provided for this image
No description has been provided for this image
No description has been provided for this image

Test d'occupation mémoire¶

Afin de mesurer l'impact de la taille de batch sur l'occupation mémoire et sur le throughput, la cellule suivante permet de soumettre plusieurs jobs avec des tailles de batch croissantes. Dans les cas où la mémoire est saturée et dépasse la capacité du GPU, le système renverra une erreur CUDA Out of Memory.

Soumission du job. Attention vous sollicitez les noeuds de calcul à ce moment-là.

Pour soumettre le job, veuillez basculer la cellule suivante du mode Raw NBConvert au mode Code.

In [12]:
n_gpu = 4
batch_sizes = [1, 2, 4, 8]
command = [f'fsdp.py --batch-size {batch_size} --num-workers 2 --seq-len 512 --test'
          for batch_size in batch_sizes]
jobids = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
                   account=account, time_max='00:10:00', constraint='a100', qos="qos_gpu_a100-dev")
print(f'jobids = {jobids}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task
Submitted batch job 335115
Submitted batch job 335120
Submitted batch job 335125
Submitted batch job 335132
jobids = ['335115', '335120', '335125', '335132']
In [13]:
display_slurm_queue(name)
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
            335115    gpu_p5 NATHOUNE  cfor032  R       0:47      1 jean-zay-iam31
            335120    gpu_p5 NATHOUNE  cfor032  R       0:47      1 jean-zay-iam31
            335125    gpu_p5 NATHOUNE  cfor032  R       0:47      1 jean-zay-iam01

 Done!
In [14]:
GPU_underthehood(jobids)
Batch size per GPU: 4 Max GPU Memory Allocated: 72.08 GB, Troughput: 10.616 images/second
Batch size per GPU: 8 Max GPU Memory Allocated: 72.32 GB, Troughput: 18.526 images/second
Batch size per GPU: 16 Max GPU Memory Allocated: 72.81 GB, Troughput: 28.875 images/second
Batch size per GPU: 32 CUDA out of memory
Memory occupancy by Model part : 71.835 +/- 0.003 GB

Passage au Fully Sharded Data Parallelism¶

TODO: Remplacez le Distributed Data Parallelism par le Fully Sharded Data Parallelism. Ce n'est qu'un simple wrapper et demande peu de modifications. Indice : ctrl-F de "#### Distribute the Model" pour repérer l'endroit où faire ça.

Important: Dans la documentation, vous pouvez voir un paramètre auto_wrap_policy, il est abordé dans la section suivante donc laissez le non précisé pour l'instant.

FSDP est un wrapper très haut niveau qui fait toutes les communications de manière cachée pour faciliter son utilisation. C'est son grand avantage par rapport à DeepSpeed.

fsdp

In [18]:
n_gpu = 4
command = f'fsdp.py --batch-size 4 --num-workers 2 --seq-len 512 --test --nccl-profile'
jobid = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
                   account=account, time_max='00:10:00')
print(f'jobid = {jobid}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task
Submitted batch job 336009
jobid = ['336009']
In [19]:
display_slurm_queue(name)
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
            336009    gpu_p5 NATHOUNE  cfor032  R       0:44      1 jean-zay-iam31

 Done!
In [20]:
#jobid = ['902284']
In [21]:
controle_technique(jobid)
Train throughput: 32.89 images/second
GPU throughput: 32.91 images/second
epoch time: 175.12 seconds
-----------
training step time average (fwd/bkwd on GPU): 0.486194 sec (60.1%/33.0%) +/- 0.065956
loading step time average (IO + CPU to GPU transfer): 0.000251 sec +/- 0.000032

Click here to display the log file

In [22]:
pipe_memory(jobid)
No description has been provided for this image
In [24]:
comm_profiler(jobid)
No description has been provided for this image
No description has been provided for this image

Test d'occupation mémoire¶

Afin de mesurer l'impact de la taille de batch sur l'occupation mémoire et sur le throughput, la cellule suivante permet de soumettre plusieurs jobs avec des tailles de batch croissantes. Dans les cas où la mémoire est saturée et dépasse la capacité du GPU, le système renverra une erreur CUDA Out of Memory.

Soumission du job. Attention vous sollicitez les noeuds de calcul à ce moment-là.

Pour soumettre le job, veuillez basculer la cellule suivante du mode Raw NBConvert au mode Code.

In [25]:
n_gpu = 4
batch_sizes = [2, 4, 8, 16]
command = [f'fsdp.py --batch-size {batch_size} --num-workers 2 --seq-len 512 --test'
          for batch_size in batch_sizes]
jobids = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
                   account=account, time_max='00:10:00', constraint='a100', qos="qos_gpu_a100-dev")
print(f'jobids = {jobids}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task
Submitted batch job 336026
Submitted batch job 336027
Submitted batch job 336028
Submitted batch job 336029
jobids = ['336026', '336027', '336028', '336029']
In [26]:
#jobids = ['299137', '299138', '299139', '299140', '299141']
In [27]:
display_slurm_queue(name)
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
            336028    gpu_p5 NATHOUNE  cfor032 CG       1:10      1 jean-zay-iam01

 Done!
In [28]:
GPU_underthehood(jobids)
Batch size per GPU: 8 Max GPU Memory Allocated: 45.39 GB, Troughput: 22.542 images/second
Batch size per GPU: 16 Max GPU Memory Allocated: 45.88 GB, Troughput: 32.915 images/second
Batch size per GPU: 32 Max GPU Memory Allocated: 53.17 GB, Troughput: 42.554 images/second
Batch size per GPU: 64 CUDA out of memory
Memory occupancy by Model part : 41.743 +/- 3.158 GB

Contrôle technique de la configuration optimale¶

In [29]:
controle_technique([jobids[-2]])
Train throughput: 42.51 images/second
GPU throughput: 42.55 images/second
epoch time: 135.50 seconds
-----------
training step time average (fwd/bkwd on GPU): 0.751985 sec (54.4%/40.4%) +/- 0.078437
loading step time average (IO + CPU to GPU transfer): 0.000791 sec +/- 0.000241

Click here to display the log file


Optimisation de la FSDP en l'appliquant récursivement¶

En utilisant la FSDP sans l'argument auto_wrap_policy, alors tous les paramètres du modèle sont transmis en une seule fois plutôt que quand ils sont nécessaires. On pourrait imaginer transmettre d'abord les paramètres des premières couches (dont on a besoin rapidement), puis ensuite les autres.

Le paramètre auto_wrap_policy permet de faire ça. Il est puissant car il permet de maximiser le recouvrement entre les calculs et les communications et minimise la mémoire nécessaire.

TODO: Ajoutez le paramètre auto_wrap_policy. Indice : Pour les transformers, il est de coutume de wrapper chaque block de transformer indépendamment (signature de la méthode transformer_auto_wrap_policy ici).

In [38]:
n_gpu = 4
command = f'fsdp.py --batch-size 4 --num-workers 2 --seq-len 512 --test --nccl-profile'
jobid = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
                   account=account, time_max='00:10:00', qos="qos_gpu_a100-dev")
print(f'jobid = {jobid}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task
Submitted batch job 336276
jobid = ['336276']
In [39]:
display_slurm_queue(name)
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
            336276    gpu_p5 NATHOUNE  cfor032  R       0:45      1 jean-zay-iam31

 Done!
In [40]:
#jobid = ['902284']
In [41]:
controle_technique(jobid)
Train throughput: 34.64 images/second
GPU throughput: 34.65 images/second
epoch time: 166.30 seconds
-----------
training step time average (fwd/bkwd on GPU): 0.461701 sec (43.9%/53.1%) +/- 0.028161
loading step time average (IO + CPU to GPU transfer): 0.000237 sec +/- 0.000031

Click here to display the log file

In [45]:
comm_profiler(jobid, n_display=100)
No description has been provided for this image
No description has been provided for this image
No description has been provided for this image

Vous pouvez consulter les logs de l'entraînement. On affiche le modèle, et on peut par conséquent voir l'impact de l'argument auto_wrap_policy.

Test d'occupation mémoire¶

Afin de mesurer l'impact de la taille de batch sur l'occupation mémoire et sur le throughput, la cellule suivante permet de soumettre plusieurs jobs avec des tailles de batch croissantes. Dans les cas où la mémoire est saturée et dépasse la capacité du GPU, le système renverra une erreur CUDA Out of Memory.

Soumission du job. Attention vous sollicitez les noeuds de calcul à ce moment-là.

Pour soumettre le job, veuillez basculer la cellule suivante du mode Raw NBConvert au mode Code.

In [46]:
n_gpu = 4
batch_sizes = [4, 8, 12, 16, 32]
command = [f'fsdp.py --batch-size {batch_size} --num-workers 2 --seq-len 512 --test'
          for batch_size in batch_sizes]
jobids = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
                   account=account, time_max='00:10:00', constraint='a100', qos="qos_gpu_a100-dev")
print(f'jobids = {jobids}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task
Submitted batch job 336313
Submitted batch job 336314
Submitted batch job 336315
Submitted batch job 336316
Submitted batch job 336318
jobids = ['336313', '336314', '336315', '336316', '336318']
In [47]:
#jobids = ['300073', '300074', '300075', '300076', '300077']
In [48]:
display_slurm_queue(name)
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
            336316    gpu_p5 NATHOUNE  cfor032  R       2:18      1 jean-zay-iam02

 Done!
In [49]:
GPU_underthehood(jobids)
Batch size per GPU: 16 Max GPU Memory Allocated: 29.56 GB, Troughput: 34.349 images/second
Batch size per GPU: 32 Max GPU Memory Allocated: 42.67 GB, Troughput: 43.954 images/second
Batch size per GPU: 48 Max GPU Memory Allocated: 55.78 GB, Troughput: 48.743 images/second
Batch size per GPU: 64 Max GPU Memory Allocated: 68.89 GB, Troughput: 29.467 images/second
Batch size per GPU: 128 CUDA out of memory
Memory occupancy by Model part : 26.281 +/- 9.830 GB

Contrôle technique de la configuration optimale¶

In [50]:
controle_technique(jobids[-3])
Train throughput: 48.73 images/second
GPU throughput: 48.74 images/second
epoch time: 118.20 seconds
-----------
training step time average (fwd/bkwd on GPU): 0.984761 sec (41.3%/56.5%) +/- 0.043922
loading step time average (IO + CPU to GPU transfer): 0.000269 sec +/- 0.000037

Click here to display the log file


Bonus : torch.compile par dessus FSDP¶

TODO: Appliquez la compilation par PyTorch de votre modèle.

Indice: ctrl-F de "#### JIT" pour trouver où faire ça.

In [55]:
n_gpu = 4
command = f'fsdp.py --batch-size 16 --num-workers 2 --seq-len 512 --test --compile --nccl-profile'
jobid = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
                   account=account, time_max='00:10:00', qos="qos_gpu_a100-dev")
print(f'jobid = {jobid}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task
Submitted batch job 336937
jobid = ['336937']
In [56]:
#jobid = ['902284']
In [57]:
display_slurm_queue(name)
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
            336937    gpu_p5 NATHOUNE  cfor032  R       2:37      1 jean-zay-iam02

 Done!
In [58]:
controle_technique(jobid)
Train throughput: 64.93 images/second
GPU throughput: 64.95 images/second
epoch time: 88.71 seconds
-----------
training step time average (fwd/bkwd on GPU): 0.985422 sec (40.5%/57.3%) +/- 0.042639
loading step time average (IO + CPU to GPU transfer): 0.000297 sec +/- 0.000049

Click here to display the log file

In [64]:
comm_profiler(jobid, n_display=100)
No description has been provided for this image
No description has been provided for this image
No description has been provided for this image

Test d'occupation mémoire¶

Afin de mesurer l'impact de la taille de batch sur l'occupation mémoire et sur le throughput, la cellule suivante permet de soumettre plusieurs jobs avec des tailles de batch croissantes. Dans les cas où la mémoire est saturée et dépasse la capacité du GPU, le système renverra une erreur CUDA Out of Memory.

Soumission du job. Attention vous sollicitez les noeuds de calcul à ce moment-là.

Pour soumettre le job, veuillez basculer la cellule suivante du mode Raw NBConvert au mode Code.

In [59]:
n_gpu = 4
batch_sizes = [8, 12, 16, 20, 24]
command = [f'fsdp.py --batch-size {batch_size} --num-workers 2 --seq-len 512 --test --compile'
          for batch_size in batch_sizes]
jobids = gpu_jobs_submitter(command, n_gpu, MODULE, name=name,
                   account=account, time_max='00:10:00', constraint='a100', qos="qos_gpu_a100-dev")
print(f'jobids = {jobids}')
batch job 0: 4 GPUs distributed on 1 nodes with 4 tasks / 4 gpus per node and 8 cpus per task
Submitted batch job 336987
Submitted batch job 336989
Submitted batch job 336991
Submitted batch job 336994
Submitted batch job 336995
jobids = ['336987', '336989', '336991', '336994', '336995']
In [60]:
#jobids = ['300403', '300404', '300406', '300407', '300408', '300409']
In [61]:
display_slurm_queue(name)
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
            336994    gpu_p5 NATHOUNE  cfor032  R       2:48      1 jean-zay-iam31

 Done!
In [62]:
GPU_underthehood(jobids)
Batch size per GPU: 32 Max GPU Memory Allocated: 35.62 GB, Troughput: 55.736 images/second
Batch size per GPU: 48 Max GPU Memory Allocated: 45.21 GB, Troughput: 61.295 images/second
Batch size per GPU: 64 Max GPU Memory Allocated: 54.80 GB, Troughput: 64.816 images/second
Batch size per GPU: 80 Max GPU Memory Allocated: 64.38 GB, Troughput: 76.294 images/second
Batch size per GPU: 96 CUDA out of memory
Memory occupancy by Model part : 36.420 +/- 5.592 GB

Contrôle technique de la configuration optimale¶

In [63]:
controle_technique(jobids[-2])
Train throughput: 69.96 images/second
GPU throughput: 76.29 images/second
epoch time: 82.33 seconds
-----------
training step time average (fwd/bkwd on GPU): 1.048578 sec (36.0%/64.8%) +/- 0.065510
loading step time average (IO + CPU to GPU transfer): 0.094912 sec +/- 0.047183

Click here to display the log file

torch.compile est encore très nouveau et il peut arriver qu'un modèle ne puisse pas être converti. Plusieurs backends sont disponibles (voir documentation officielle). Dans les cas des modèles les plus exotiques, la compilation peut tout simplement échoué. C'est tellement bas niveau qu'il est bien possible qu'on ne puisse rien y faire, c'est juste lié au fait que torch.compile est relativement nouveau. À garder à l'esprit cependant, car cela peut augmenter de 50%, voire parfois 100% le throughput de votre modèle.

Commentaires